//=============================================================================
// FILE:
//    CmgDudePass.cpp
//
// DESCRIPTION:
//    Visits all functions in a module, prints their names and the number of
//    arguments via stderr. Strictly speaking, this is an analysis pass (i.e.
//    the functions are not modified). However, in order to keep things simple
//    there's no 'print' method here (every analysis pass should implement it).
//
// USAGE:
//    1. Legacy PM
//      opt -load libCmgDudePass.so -LegacyCmgDudePass -disable-output
//      <input-llvm-file>
//    2. New PM
//      opt -load-pass-plugin=libCmgDudePass.so -passes="CmgDudePass" -disable-output
//      <input-llvm-file>
//    3. Clang Embedding
//      clang/clang++ -Xclang -load -Xclang libCmgDudePass.so -emit-llvm -S
//      <input-source-file>
//
// License: MIT
//=============================================================================
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/Analysis/TypeMetadataUtils.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Value.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/Transforms/IPO/WholeProgramDevirt.h"

#include <Database.h>

#include <OUT.h>

using namespace llvm;
using namespace llvm::wholeprogramdevirt;
using namespace cmg;

namespace {
// A virtual call site. VTable is the loaded virtual table pointer, and CS
// is the indirect virtual call.
struct VirtualCallSite {
  Value *VTable = nullptr;
  CallBase &CB;
};

// Call site information collected for a specific VTableSlot and possibly a
// list of constant integer arguments. The grouping by arguments is handled by
// the VTableSlotInfo class.
struct CallSiteInfo {
  /// The set of call sites for this slot. Used during regular LTO and the
  /// import phase of ThinLTO (as well as the export phase of ThinLTO for any
  /// call sites that appear in the merged module itself); in each of these
  /// cases we are directly operating on the call sites at the IR level.
  std::vector<VirtualCallSite> CallSites;
};

// A slot in a set of virtual tables. The TypeID identifies the set of
// virtual tables, and the ByteOffset is the offset in bytes from the
// address point to the virtual function pointer.
struct VTableSlot {
  Metadata *TypeID;
  uint64_t ByteOffset;
}; // namespace VTableSlot

// Call site information collected for a specific VTableSlot.
struct VTableSlotInfo {
  // The set of call sites which do not have all constant integer arguments
  // (excluding "this").
  CallSiteInfo CSInfo;

  // The set of call sites with all constant integer arguments (excluding
  // "this"), grouped by argument list.
  std::map<std::vector<uint64_t>, CallSiteInfo> ConstCSInfo;

  void addCallSite(Value *VTable, CallBase &CB) {
    CSInfo.CallSites.push_back({VTable, CB});
  }

private:
  CallSiteInfo &findCallSiteInfo(CallBase &CB) {}
};

struct CollectVirtualCandidate {
  Module &M;
  std::string &ModuleName;
  function_ref<DominatorTree &(Function &)> LookupDomTree;
  MapVector<VTableSlot, VTableSlotInfo> CallSlots;
  Connection &conn;

  CollectVirtualCandidate(
      Module &M, function_ref<DominatorTree &(Function &)> LookupDomTree,
      string &ModuleName, Connection &conn)
      : M(M), LookupDomTree(LookupDomTree), ModuleName(ModuleName), conn(conn) {
  }

  void
  scanTypeTestUsers(Function *TypeTestFunc,
                    DenseMap<Metadata *, std::set<TypeMemberInfo>> &TypeIdMap) {
    for (auto I = TypeTestFunc->use_begin(); I != TypeTestFunc->use_end();
         I++) {
      auto CI = dyn_cast<CallInst>(I->getUser());
      if (!CI)
        continue;
      SmallVector<DevirtCallSite, 1> DevirtCalls;
      SmallVector<CallInst *, 1> Assumes;
      auto &DT = LookupDomTree(*CI->getFunction());
      findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI, DT);

      Metadata *TypeId =
          cast<MetadataAsValue>(CI->getArgOperand(1))->getMetadata();

      if (!Assumes.empty()) {
        Value *Ptr = CI->getArgOperand(0)->stripPointerCasts();
        for (DevirtCallSite Call : DevirtCalls)
          CallSlots[{TypeId, Call.Offset}].addCallSite(Ptr, Call.CB);
      }
    }
  }

  void scanTypeCheckedLoadUsers(Function *TypeCheckedLoadFunc) {
    for (auto I = TypeCheckedLoadFunc->use_begin();
         I != TypeCheckedLoadFunc->use_end(); I++) {
      auto CI = dyn_cast<CallInst>(I->getUser());
      if (!CI)
        continue;

      Value *Ptr = CI->getArgOperand(0);
      Value *Offset = CI->getArgOperand(1);
      Value *TypeIdValue = CI->getArgOperand(2);
      Metadata *TypeId = cast<MetadataAsValue>(TypeIdValue)->getMetadata();

      SmallVector<DevirtCallSite, 1> DevirtCalls;
      SmallVector<Instruction *, 1> LoadedPtrs;
      SmallVector<Instruction *, 1> Preds;
      bool HasNonCallUses = false;
      auto &DT = LookupDomTree(*CI->getFunction());
      findDevirtualizableCallsForTypeCheckedLoad(DevirtCalls, LoadedPtrs, Preds,
                                                 HasNonCallUses, CI, DT);

      for (DevirtCallSite Call : DevirtCalls) {
        CallSlots[{TypeId, Call.Offset}].addCallSite(Ptr, Call.CB);
      }
    }
  }

  void buildTypeIdentifierMap(
      std::vector<VTableBits> &Bits,
      DenseMap<Metadata *, std::set<TypeMemberInfo>> &TypeIdMap) {
    DenseMap<GlobalVariable *, VTableBits *> GVToBits;
    Bits.reserve(M.getGlobalList().size());
    SmallVector<MDNode *, 2> Types;
    for (GlobalVariable &GV : M.globals()) {
      Types.clear();
      GV.getMetadata(LLVMContext::MD_type, Types);
      if (GV.isDeclaration() || Types.empty())
        continue;

      VTableBits *&BitsPtr = GVToBits[&GV];
      if (!BitsPtr) {
        Bits.emplace_back();
        Bits.back().GV = &GV;
        Bits.back().ObjectSize =
            M.getDataLayout().getTypeAllocSize(GV.getInitializer()->getType());
        BitsPtr = &Bits.back();
      }

      OUT_ << ModuleName << " "
           << "VTable:\n";
      OUT_ << "\t";
      // GV.print(OUT_);
      OUT_ << GV.getName();
      OUT_ << "\n";
      OUT_ << "\tType ";
      GV.getInitializer()->getType()->print(OUT_);
      OUT_ << "\n";
      OUT_ << "\tVCallVisibility: " << GV.getVCallVisibility() << "\n";
      OUT_ << "\tIsArrayTy: " << GV.getInitializer()->getType()->isArrayTy()
           << "\n";
      unsigned int vtableSize =
          M.getDataLayout().getTypeAllocSize(GV.getInitializer()->getType());
      OUT_ << "\tVTable size: "
           << M.getDataLayout().getTypeAllocSize(GV.getInitializer()->getType())
           << "\n";
      PointerType *i8 = Type::getInt8PtrTy(M.getContext());
      unsigned int i8_size = M.getDataLayout().getTypeAllocSize(i8);
      OUT_ << "i8 size:" << i8_size << "\n";

      unsigned int vtableEntryNum = vtableSize / i8_size;

      for (unsigned int i = 0; i < vtableEntryNum; i++) {
        Constant *entry =
            getPointerAtOffset(GV.getInitializer(), i * i8_size, M);
        if (!entry) {
          OUT_ << "\t" << i * i8_size << " null\n";
        } else {
          auto *FN = dyn_cast<Function>(entry->stripPointerCasts());
          if (!FN) {
            OUT_ << "\t" << i * i8_size << " null\n";
          } else {
            OUT_ << "\t" << i * i8_size << " " << FN->getName() << "\n";
            string vtable = GV.getName().str();
            string virtual_function = FN->getName().str();
            conn.insertToVtables(ModuleName, vtable, i * i8_size,
                                 virtual_function);
          }
        }
      }

      for (MDNode *Type : Types) {
        auto TypeID = Type->getOperand(1).get();

        uint64_t Offset =
            cast<ConstantInt>(
                cast<ConstantAsMetadata>(Type->getOperand(0))->getValue())
                ->getZExtValue();

        TypeIdMap[TypeID].insert({BitsPtr, Offset});
      }
    }
  }

  bool
  tryFindVirtualCallTargets(std::vector<VirtualCallTarget> &TargetsForSlot,
                            const std::set<TypeMemberInfo> &TypeMemberInfos,
                            uint64_t ByteOffset) {

    return true;
  }

  bool run() {
    OUT_ << "Module: " << M.getName() << "\n";

    Function *TypeTestFunc =
        M.getFunction(Intrinsic::getName(Intrinsic::type_test));
    Function *TypeCheckedLoadFunc =
        M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load));
    Function *AssumeFunc = M.getFunction(Intrinsic::getName(Intrinsic::assume));

    // build type metadata into a map for easy lookup
    std::vector<VTableBits> Bits;
    DenseMap<Metadata *, std::set<TypeMemberInfo>> TypeIdMap;
    buildTypeIdentifierMap(Bits, TypeIdMap);
    for (auto &p : TypeIdMap) {
      Metadata *TypeId = p.first;
      OUT_ << ModuleName << " MetaType: " << cast<MDString>(TypeId)->getString()
           << "\n";
      for (const TypeMemberInfo &tmi : p.second) {
        OUT_ << "\tGV: " << tmi.Bits->GV->getName() << " " << tmi.Offset
             << "\n";
        string type_metadata = cast<MDString>(TypeId)->getString().str();
        string vtable = tmi.Bits->GV->getName().str();
        conn.insertToTypeMetadata(type_metadata, vtable, tmi.Offset);
      }
    }

    if (TypeTestFunc && AssumeFunc) {
      scanTypeTestUsers(TypeTestFunc, TypeIdMap);
    }

    if (TypeCheckedLoadFunc) {
      scanTypeCheckedLoadUsers(TypeCheckedLoadFunc);
    }

    for (auto P : CallSlots) {
      OUT_ << ModuleName << " VirtualCall "
           << dyn_cast<MDString>(P.first.TypeID)->getString() << " "
           << P.first.ByteOffset << "\n";
      //  OUT_<<"\tCallInst "<<
      string type_metadata =
          dyn_cast<MDString>(P.first.TypeID)->getString().str();

      for (auto VCS : P.second.CSInfo.CallSites) {
        OUT_ << "\tVTable ";
        //  << VCS.VTable->getName() << "\n";
        VCS.VTable->print(OUT_);
        OUT_ << "\n";
        OUT_ << "\tCallSite ";
        VCS.CB.print(OUT_);
        OUT_ << "\n";
        Function *FN = VCS.CB.getParent()->getParent();
        string user_symbol = FN->getName().str();
        conn.insertToVitualCall(ModuleName, user_symbol, type_metadata,
                                P.first.ByteOffset);
      }
    }
    return true;
  }
};

} // namespace

namespace llvm {

template <> struct DenseMapInfo<VTableSlot> {
  static VTableSlot getEmptyKey() {
    return {DenseMapInfo<Metadata *>::getEmptyKey(),
            DenseMapInfo<uint64_t>::getEmptyKey()};
  }
  static VTableSlot getTombstoneKey() {
    return {DenseMapInfo<Metadata *>::getTombstoneKey(),
            DenseMapInfo<uint64_t>::getTombstoneKey()};
  }
  static unsigned getHashValue(const VTableSlot &I) {
    return DenseMapInfo<Metadata *>::getHashValue(I.TypeID) ^
           DenseMapInfo<uint64_t>::getHashValue(I.ByteOffset);
  }
  static bool isEqual(const VTableSlot &LHS, const VTableSlot &RHS) {
    return LHS.TypeID == RHS.TypeID && LHS.ByteOffset == RHS.ByteOffset;
  }
};
} // namespace llvm

//-----------------------------------------------------------------------------
// CmgDudePass implementation
//-----------------------------------------------------------------------------
// No need to expose the internals of the pass to the outside world - keep
// everything in an anonymous namespace.
namespace {

void FindDirectGlobalValueUser(Value *usee, vector<GlobalValue *> &rst) {
  for (auto user : usee->users()) {
    if (Function *FN = dyn_cast<Function>(user)) {
      rst.push_back(FN);
    } else if (Instruction *INST = dyn_cast<Instruction>(user)) {
      Function *FN = INST->getFunction();
      rst.push_back(FN);
    } else if (GlobalVariable *GV = dyn_cast<GlobalVariable>(user)) {
      rst.push_back(GV);
    } else {

      if (user->user_empty())
        return;

      FindDirectGlobalValueUser(user, rst);
    }
  }
}

// New PM implementation
struct CmgDudePass : PassInfoMixin<CmgDudePass> {
  // Main entry point, takes IR unit to run the pass on (&F) and the
  // corresponding pass manager (to be queried if need be)
  // PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
  //   visitor(F);
  //   return PreservedAnalyses::all();
  // }
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {

    return PreservedAnalyses::all();
  }
};

// Legacy PM implementation
struct LegacyCmgDudePass : public ModulePass {
  // Connection conn;
  string ModuleName;
  static char ID;
  LegacyCmgDudePass() : ModulePass(ID) {}
  // Main entry point - the name conveys what unit of IR this is to be run on.
  bool runOnModule(Module &M) override {
    auto LookForDominatorTree = [&](Function &F) -> DominatorTree & {
      return this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
    };
    ModuleName = M.getSourceFileName();

    llvm::SmallVector<char, 50> current_working_path;
    llvm::sys::fs::current_path(current_working_path);
    ModuleName = string(current_working_path.data()) + "/" + ModuleName;
    // llvm::outs() << "Module name: " << ModuleName << "\n";

    // if (M.Directory) {
    //   StringRef dir = M.Directory->getName();
    //   ModuleName = (dir + M.getName()).str();
    // } else {
    //   ModuleName = M.getName().str();
    // }
    Connection conn(ModuleName);

    CollectVirtualCandidate{M, LookForDominatorTree, ModuleName, conn}.run();

    for (auto &global : M.global_values()) {
      if (global.isDSOLocal() && global.hasExactDefinition())
        continue;

      if (global.isDeclarationForLinker()) {
        if (!global.user_empty()) {
          OUT_ << "Symbol dependency: ";
          OUT_ << global.getName();
          OUT_ << "\n";
          vector<GlobalValue *> globalUsers;
          FindDirectGlobalValueUser(&global, globalUsers);
          if (globalUsers.empty()) {
            OUT_ << "\tnot used by Global Values;\n";
          } else {
            for (Value *user : globalUsers) {
              if (Function *FN = dyn_cast<Function>(user)) {
                OUT_ << "\tused by Function: " << FN->getName() << "\n";
                string user_symbol = FN->getName().str();
                string used_symbol = global.getName().str();
                conn.insertToSymbolUse(ModuleName, user_symbol, used_symbol);
              } else if (GlobalVariable *GV = dyn_cast<GlobalVariable>(user)) {
                OUT_ << "\tused by GlobalVariable: " << GV->getName() << "\n";
                string user_symbol = GV->getName().str();
                string used_symbol = global.getName().str();
                conn.insertToSymbolUse(ModuleName, user_symbol, used_symbol);
              }
            }
          }
        }
      }
    }
    for (auto &global : M.global_values()) {
      if (global.hasExactDefinition() && global.hasExternalLinkage()) {
        OUT_ << "Symbol definition ";
        if (global.getValueType()->isFunctionTy()) {
          OUT_ << "Func : ";
          OUT_ << global.getName();
          OUT_ << "\n";
          string symbol = global.getName().str();
          conn.insertToSymbolDefinition(ModuleName, symbol, 1);
        } else {
          OUT_ << "GV : ";
          OUT_ << global.getName();
          OUT_ << "\n";
          string symbol = global.getName().str();
          conn.insertToSymbolDefinition(ModuleName, symbol, 0);
        }
      }
    }

    // Doesn't modify the input unit of IR, hence 'false'
    return false;
  }
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<DominatorTreeWrapperPass>();
  }
};
} // namespace

//-----------------------------------------------------------------------------
// New PM Registration
// opt -load-pass-plugin <path/to/libCmgDudePass.so> -passes="CmgDudePass" <.s/.ll/.bc>
//-----------------------------------------------------------------------------
llvm::PassPluginLibraryInfo getCmgDudePassPluginInfo() {
  return {LLVM_PLUGIN_API_VERSION, "CmgDudePass", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager &MAM,
                   ArrayRef<PassBuilder::PipelineElement>) {
                  if (Name == "CmgDudePass") {
                    MAM.addPass(CmgDudePass());
                    return true;
                  }
                  return false;
                });
          }};
}

// This is the core interface for pass plugins. It guarantees that 'opt' will
// be able to recognize CmgDudePass when added to the pass pipeline on the
// command line, i.e. via '-passes=CmgDudePass'
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return getCmgDudePassPluginInfo();
}

//-----------------------------------------------------------------------------
// Legacy PM Registration
// opt -load <path/to/libCmgDudePass.so> --LegacyCmgDudePass -analyze <.s/.ll/.bc>
//-----------------------------------------------------------------------------
// The address of this variable is used to uniquely identify the pass. The
// actual value doesn't matter.
char LegacyCmgDudePass::ID = 0;

// This is the core interface for pass plugins. It guarantees that 'opt' will
// recognize LegacyCmgDudePass when added to the pass pipeline on the command
// line, i.e.  via '--LegacyCmgDudePass'
static RegisterPass<LegacyCmgDudePass>
    X("LegacyCmgDudePass", "Cmg Pass",
      true, // This pass doesn't modify the CFG => true
      false // This pass is not a pure analysis pass => false
    );

//-----------------------------------------------------------------------------
// Clang Registration
// clang -Xclang -load -Xclang -load <.so>
//-----------------------------------------------------------------------------
static RegisterStandardPasses RegisterCmgDudePass(
    PassManagerBuilder::EP_ModuleOptimizerEarly,
    [](const PassManagerBuilder &, legacy::PassManagerBase &PM) {
      PM.add(new LegacyCmgDudePass());
    });